package com.chatplus.application.processor.channel.chat.baidu;

import baidumodel.constant.ChatRoleConstant;
import baidumodel.entity.chat.BaiduChatMessage;
import baidumodel.entity.chat.ErnieBot4Param;
import baidumodel.entity.chat.ErnieBotTurboParam;
import baidumodel.service.BaiduService;
import baidumodel.util.TokenUtil;
import com.chatplus.application.common.logging.SouthernQuietLogger;
import com.chatplus.application.common.logging.SouthernQuietLoggerFactory;
import com.chatplus.application.domain.dto.ChatConfigDto;
import com.chatplus.application.domain.entity.basedata.ApiKeyEntity;
import com.chatplus.application.processor.AIChannelServiceProvider;
import com.chatplus.application.processor.channel.chat.IChatAiProcessorService;
import com.chatplus.application.processor.constant.AIConstants;
import com.chatplus.application.processor.constant.AiPlatformConstants;
import com.chatplus.application.processor.handler.dto.ChatRecordMessage;
import com.chatplus.application.processor.listener.BaiduEventSourceListener;
import com.chatplus.application.service.basedata.ApiKeyService;
import com.chatplus.application.service.basedata.ConfigService;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;


/**
 * 百度 文心一言 机器人处理器
 *
 * @author chj
 * @date 2023/12/27
 **/
@Service(value = AIChannelServiceProvider.SERVICE_NAME_PRE + "BaiDu")
public class BaiDuAiProcessor extends IChatAiProcessorService {
    private static final SouthernQuietLogger LOGGER = SouthernQuietLoggerFactory.getLogger(BaiDuAiProcessor.class);
    private final ConfigService configService;
    private final ApiKeyService apiKeyService;

    public BaiDuAiProcessor(ConfigService configService
            , ApiKeyService apiKeyService) {
        this.configService = configService;
        this.apiKeyService = apiKeyService;
    }

    @Override
    public void process(List<ChatRecordMessage> messageContext) {
        ChatConfigDto chatConfigDto = configService.getChatConfig();
        ChatConfigDto.BaiduBean baiduBean = chatConfigDto.getBaidu();
        if (baiduBean == null) {
            replyErrorMessage("加载系统配置失败，连接已关闭！！！");
            return;
        }
        List<ApiKeyEntity> apiKeys = apiKeyService.getApiKeys(getChannel(), "chat");
        // // 更新 API KEY 的最后使用时间
        if (CollectionUtils.isEmpty(apiKeys)) {
            replyErrorMessage("抱歉😔😔😔，系统已经没有可用的 API KEY，请联系管理员！");
            replyErrorMessage(AIConstants.ERROR_MSG);
            replyErrorMessage(AIConstants.ERR_IMG);
            return;
        }
        String prompt = messageContext.stream()
                .filter(msg -> msg.getReply() == null)
                .map(ChatRecordMessage::getPrompt)
                .collect(Collectors.joining());
        // TODO 这个后续可以考虑弄成多个配置，做轮询
        ApiKeyEntity apiKeyEntity = apiKeys.get(0);
        String apiKey = apiKeyEntity.getValue().split("\\|")[0];
        String secretKey = apiKeyEntity.getValue().split("\\|")[1];
        String modelName = getSession().getAttributes().get("modelName").toString();
        BaiduService baiduService = new BaiduService(apiKey, secretKey);
        String token = TokenUtil.getToken(apiKey, secretKey);
        // （1）messages成员不能为空，1个成员表示单轮对话，多个成员表示多轮对话
        // （2）最后一个message为当前请求的信息，前面的message为历史对话信息
        // （3）必须为奇数个成员，成员中message的role必须依次为user、assistant
        //  (4）最后一个message的content长度（即此轮对话的问题）不能超过4800个字符，且不能超过2000 tokens
        // （5）如果messages中content总长度大于4800个字符或2000 tokens，系统会依次遗忘最早的历史会话，直到content的总长度不超过4800个字符且不超过2000 tokens
        List<BaiduChatMessage> messages = new ArrayList<>();
        messageContext.forEach(msg -> {
            BaiduChatMessage promptMessage = BaiduChatMessage.builder()
                    .role(ChatRoleConstant.USER)
                    .content(msg.getPrompt())
                    .build();
            messages.add(promptMessage);
            if (StringUtils.isNotEmpty(msg.getReply())) {
                BaiduChatMessage replyMessage = BaiduChatMessage.builder()
                        .role(ChatRoleConstant.ASSISTANT)
                        .content(msg.getReply())
                        .build();
                messages.add(replyMessage);
            }
        });
        switch (modelName) {
            // 文心一言3  eb-instant
            case "eb-instant":
                ErnieBotTurboParam param3 = ErnieBotTurboParam.builder()
                        .messages(messages)
                        .user_id(getSession().getAttributes().get("userId").toString())
                        .temperature(Float.parseFloat(baiduBean.getTemperature() + ""))
                        .penalty_score(1.0F)
                        .top_p(0.8F)
                        .build();

                baiduService.ernieBotTurboStream(param3, token, new BaiduEventSourceListener(getSession(), getChannel(), prompt));
                break;
            // 文心一言4  completions_pro
            case "completions_pro":
                ErnieBot4Param param4 = ErnieBot4Param.builder()
                        .messages(messages)
                        .user_id(getSession().getAttributes().get("userId").toString())
                        .temperature(Float.parseFloat(baiduBean.getTemperature() + ""))
                        .penalty_score(1.0F)
                        .top_p(0.8F)
                        .build();
                baiduService.ernieBot4Stream(param4, token, new BaiduEventSourceListener(getSession(), getChannel(), prompt));
                break;
            default:
                replyErrorMessage("未知模型");
        }
    }

    @Override
    public String getChannel() {
        return AiPlatformConstants.BAI_DU;
    }
}
